"""
手写数字识别 Web应用
使用Gradio创建Web界面，结合标准化器与KNN模型进行手写数字识别
"""

# 导入必要的库和模块
import gradio as gr
import numpy as np
import pickle
import cv2
from sklearn.preprocessing import StandardScaler  # 新增：加载标准化器

# 加载保存的KNN模型和标准化器
print("正在加载预训练的KNN模型和标准化器...")
try:
    with open('best_knn_model.pkl', 'rb') as f:
        knn_model = pickle.load(f)
    with open('scaler.pkl', 'rb') as f:
        scaler = pickle.load(f)
    print("KNN模型和标准化器加载成功！")
    print(f"模型使用的K值: {knn_model.n_neighbors}")
except FileNotFoundError:
    print("错误：未找到模型文件 'best_knn_model.pkl' 或标准化器 'scaler.pkl'")
    print("请先运行 optimal_knn.py 来训练并保存模型")
    exit(1)

def preprocess_image(image):
    """优化后的图像预处理：确保与训练数据（8x8、白底黑字、0-16灰度）一致，且经过标准化"""
    if image is None:
        return None
    
    # 1. 统一提取图像数据（兼容Gradio Sketchpad的输出格式）
    if isinstance(image, dict):
        image = image.get('image', image.get('mask', next(iter(image.values()))))
    
    # 2. 转换为numpy数组并统一数据类型
    img_array = np.array(image) if not isinstance(image, np.ndarray) else image.copy()
    if img_array.dtype != np.uint8:
        img_array = (img_array - np.min(img_array)) / (np.max(img_array) - np.min(img_array) + 1e-8)
        img_array = (img_array * 255).astype(np.uint8)
    
    # 3. 转为灰度图（强制去除RGB/Alpha通道干扰）
    if len(img_array.shape) == 3:
        if img_array.shape[2] == 4:  # 含Alpha通道（透明背景）
            img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGBA2GRAY)
            # 去除Alpha通道的透明背景（设为白色）
            alpha_channel = img_array[:, :, 3]
            img_gray[alpha_channel < 128] = 255  # 透明区域填充为白色
        else:  # RGB通道
            img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    else:
        img_gray = img_array
    
    # 4. 二值化处理（过滤浅灰色噪声，只保留深黑数字）
    _, img_binary = cv2.threshold(
        img_gray, 0, 255, 
        cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU  # 自动反转（黑字→白背景黑字）
    )
    
    # 5. 图像居中与裁剪（解决用户手绘偏移问题）
    contours, _ = cv2.findContours(img_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        x, y, w, h = cv2.boundingRect(contours[0])
        img_cropped = img_binary[y:y+h, x:x+w]
        # 调整为8x8，保持宽高比
        scale = min(8 / w, 8 / h) if w > 0 and h > 0 else 1.0
        new_w, new_h = int(w * scale), int(h * scale)
        img_resized = cv2.resize(img_cropped, (new_w, new_h), interpolation=cv2.INTER_AREA)
        # 创建8x8画布，将缩放后的数字居中放置
        img_8x8 = np.ones((8, 8), dtype=np.uint8) * 255  # 白底
        offset_x = (8 - new_w) // 2
        offset_y = (8 - new_h) // 2
        img_8x8[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = img_resized
    else:
        img_8x8 = np.ones((8, 8), dtype=np.uint8) * 255  # 无数字时返回全白
    
    # 6. 归一化到0-16 + 标准化（与训练流程完全一致）
    img_normalized = (img_8x8 / 255.0) * 16  # 黑字（0）→ 白底（16）
    img_flattened = img_normalized.flatten().reshape(1, -1)
    img_scaled = scaler.transform(img_flattened)  # 用训练好的标准化器处理
    
    return img_scaled

def predict_digit(image):
    """
    预测手写数字（优化版：严格类型检查 + 鲁棒置信度计算）
    Args:
        image: 输入图像 (PIL Image 或 numpy array)
    Returns:
        predicted_digit: 预测的数字（-1表示失败）
        confidence: 预测置信度（0-1）
    """
    try:
        processed_img = preprocess_image(image)
        if processed_img is None:
            return -1, 0.0
        
        # 1. 模型预测（严格类型检查）
        prediction = knn_model.predict(processed_img)
        if not isinstance(prediction, np.ndarray) or prediction.size == 0:
            return -1, 0.0
        predicted_digit = int(prediction[0])  # 取第一个预测结果
        
        # 2. 置信度计算（优先用predict_proba，否则用距离倒数近似）
        confidence = 0.0
        if hasattr(knn_model, 'predict_proba'):
            probs = knn_model.predict_proba(processed_img)
            if isinstance(probs, np.ndarray) and probs.size > 0:
                confidence = float(np.max(probs))
        else:
            # KNN无predict_proba时，用距离倒数近似置信度
            distances, _ = knn_model.kneighbors(processed_img)
            if distances.size > 0 and distances[0, 0] > 0:
                confidence = 1.0 / distances[0, 0]
                confidence = min(confidence, 1.0)  # 限制最大置信度为1
        
        return predicted_digit, confidence
        
    except Exception as e:
        print(f"预测错误: {e}")
        return -1, 0.0

def process_sketchpad(input_image):
    if input_image is None:
        return "请先在画板中央绘制数字", "0.0%"
    
    digit, confidence = predict_digit(input_image)
    confidence_percent = f"{confidence*100:.1f}%"
    
    if digit == -1:
        if confidence < 0.3:
            return "预测失败：数字不清晰/太小，请用深色笔画大一些", confidence_percent
        else:
            return "预测失败：无法确定数字，请调整位置后重画", confidence_percent
    else:
        return f"预测结果: {digit}", confidence_percent

def create_interface():
    """创建Gradio用户界面（优化布局与提示）"""
    description = """
    # 🎯 手写数字识别器
    使用**最优KNN模型**识别手写数字（0-9）
    **使用步骤**:
    1. 在下方画板用鼠标/手写笔绘制数字
    2. 点击「识别数字」按钮
    3. 查看右侧预测结果与置信度
    **绘制提示**:
    - 尽量在画板中央绘制，保持数字清晰
    - 数字大小建议占画板1/3 - 1/2区域
    - 推荐使用黑色画笔，笔画粗细适中
    """
    
    with gr.Blocks(title="手写数字识别", theme=gr.themes.Soft()) as demo:
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column(scale=1):
                sketchpad = gr.Sketchpad(
                    label="绘制数字区域",
                    type="numpy"
                )
                with gr.Row():
                    clear_btn = gr.Button("清空画板", variant="secondary")
                    submit_btn = gr.Button("识别数字", variant="primary")
            
            with gr.Column(scale=1):
                result_output = gr.Textbox(
                    label="预测结果",
                    placeholder="预测结果将显示在这里...",
                    interactive=False
                )
                confidence_output = gr.Textbox(
                    label="置信度",
                    placeholder="置信度将显示在这里...",
                    interactive=False
                )
        
        # 绑定事件
        submit_btn.click(
            fn=process_sketchpad,
            inputs=sketchpad,
            outputs=[result_output, confidence_output]
        )
        clear_btn.click(
            fn=lambda: (None, "请先绘制一个数字", "0.0%"),
            inputs=None,
            outputs=[sketchpad, result_output, confidence_output]
        )
        
        with gr.Accordion("模型详情", open=False):
            gr.Markdown(f"""
            - **模型类型**: K近邻 (KNN)
            - **最优K值**: {knn_model.n_neighbors}
            - **特征处理**: 8×8像素归一化 + 标准化
            - **识别范围**: 数字 0-9
            """)
    
    return demo

if __name__ == "__main__":
    print("正在启动Gradio Web应用...")
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",  # 允许外部访问
        server_port=7862,
        share=True,
        inbrowser=True
    )